from .Vector import Vector
from .Matrix import Matrix

class LinearSystem:
    def __init__(self, A, b):
        assert A.row_num() == len(b), "A row must == b len"
        self._m = A.row_num()
        self._n = A.col_num()
        assert self._m == self._n # TODO: no this restriction

        self.Ab = [ Vector(A.row_vector(i).underlying_list() + [b[i]])
                    for i in range(self._m)]

    def gauss_jordan_elimination(self):
        self._forward()
        self._backward()

    def _forward(self):
        for i in range(self._m):
            # 交换最大行，提高数值计算的精度
            max_row = self._max_row(i, self._m)
            self.Ab[i], self.Ab[max_row] = self.Ab[max_row], self.Ab[i]

            #将主元归一，Ab[i][i]为主元
            self.Ab[i] = self.Ab[i] / self.Ab[i][i]

            #同一列下方元素置0
            for j in range(i + 1, self._m):
                self.Ab[j] = self.Ab[j] - self.Ab[i] * self.Ab[j][i]

    def _backward(self):
        for i in range(self._m - 1, -1, -1):
            for j in range(i - 1, -1, -1):
                self.Ab[j] = self.Ab[j] - self.Ab[i] * self.Ab[j][i]

    def _max_row(self, index, n):
        best, ret = self.Ab[index][index], index
        for i in range(index + 1, n):
            if self.Ab[i][index] > best:
                best, ret = self.Ab[i][index], i
        return ret

    def fancy_print(self):
        for i in range(self._m):
            print(" ".join(str(self.Ab[i][j]) for j in range(self._n)), end = " ")
            print("|", self.Ab[i][-1])